# -*- encoding: utf-8 -*-
'''
@File    :   mae_dagmm.py
@Time    :   2021/12/1 16:47
@Author  :   ZhangChaoYang
@Desc    :   先用MAE对原始数据进行降维预处理，再用DAGMM进行异常检测。
'''
import os
import sys

sys.path.insert(0, os.getcwd())
import tensorflow as tf
import numpy as np
from util.err_analyze import fit_err_percentage
from anomaly_detection.gmm import GMM
from anomaly_detection.dagmm import DAGMM, FLAGS
from models import losses
from util.work_flow import preprocess, gen_outfile, save_model, load_model, train_test_split
import tensorflow_addons as tfa
import pickle
from absl import app


def main(argv):
    corpus = FLAGS.corpus
    data_dim = FLAGS.data_dim
    data_trans = FLAGS.data_trans

    learning_rate = 1e-4
    batch_size = 256
    epochs = 100  # paderborn数据量大，轮次可以少一些
    n_component = 4
    model_name = "mae_dagmm"

    mae = tf.keras.models.load_model(os.path.join("data", "model", data_dim, corpus, data_trans, "mae"),
                                     compile=False)
    encoder = mae.encoder
    encoder.down_stream = True

    normal_data_files, anomaly_data_file, train_history_file, check_file, model_file, scaler_file = gen_outfile(
        data_dim,
        corpus,
        data_trans,
        model_name)

    X, ano_X = preprocess(normal_data_files, anomaly_data_file, data_dim)
    X = encoder.predict(X, batch_size=10000)  # predict()返回的是numpy，不是Tensor
    ano_X = encoder.predict(ano_X, batch_size=10000)  # 输入数据很大时，用predict指定batch_size
    X_train, X_test = train_test_split(X)

    gmm = GMM()
    model = DAGMM(input_shape=(X.shape[-2], X.shape[-1]), hidden_dims=[128, 32], estimater_dims=[10, n_component])
    optimizer = tfa.optimizers.AdamW(learning_rate=learning_rate, weight_decay=1E-4)
    model.train(X_train, X_test, batch_size=batch_size, epochs=epochs, optimizer=optimizer,
                chart_file=train_history_file)
    save_model(model, model_file)
    with open(os.path.join(model_file, "opt"), "wb") as fout:
        pickle.dump(optimizer.get_config(), fout, protocol=pickle.HIGHEST_PROTOCOL)

    print("正常样本")
    x_hat, gamma, normal_z = model(X)
    fit_err_percentage(X, x_hat, losses.square_loss)  # 编码器在正常样本上的拟合误差
    gmm.fit(normal_z, gamma)
    # DAGMM建议按energy loss划定阈值,以些来区分正常和异常样本
    energy = gmm.energy(normal_z)
    print("energy loss")
    sl = sorted(energy.numpy().tolist())
    for p in range(0, 100, 5):
        index = int(p / 100.0 * len(sl))
        print("{}%\t{:.2e}".format(p, sl[index]))
    print("100%\t{:.2e}".format(sl[-1]))
    gmm.threshold = sl[int(0.9 * len(sl))]  # 以正常样本90%分位点作为判别阈值
    print("threshold", gmm.threshold)
    gmm.save(os.path.join(model_file, "gmm"))

    model = load_model(model_file)
    gmm.load(os.path.join(model_file, "gmm"))

    import matplotlib.pyplot as plt

    print("异常样本")
    ano_x_hat, gamma, ano_z = model(ano_X)
    fit_err_percentage(ano_X, ano_x_hat, losses.square_loss)  # 编码器在异常样本上的拟合误差

    # DAGMM建议按energy loss划定阈值,以些来区分正常和异常样本
    energy = gmm.energy(ano_z)
    print("energy loss")
    ano_sl = sorted(energy.numpy().tolist())
    for p in range(0, 100, 5):
        index = int(p / 100.0 * len(ano_sl))
        print("{}%\t{:.2e}".format(p, ano_sl[index]))
    print("100%\t{:.2e}".format(ano_sl[-1]))
    print("异常样本召回率{:.2f}%".format(100 * np.count_nonzero(np.asarray(ano_sl) > gmm.threshold) / len(ano_sl)))

    if len(sl) > 10000:
        sl = np.random.choice(sl, 10000)
        sl = sorted(sl)
    if len(ano_sl) > 10000:
        ano_sl = np.random.choice(ano_sl, 10000)
        ano_sl = sorted(ano_sl)
    plt.figure(figsize=(8, 8))
    ax = plt.subplot()
    plt.suptitle("蓝色:正常样本,红色:异常样本")
    ax.scatter(range(len(sl) // 100 * 99), sl[:len(sl) // 100 * 99], color='b')  # 正常样本用蓝色
    ax.scatter(range(len(ano_sl) // 100 * 99), ano_sl[:len(ano_sl) // 100 * 99], color='r')  # 异常样本用红色
    plt.savefig(check_file, format="png")

    # from sklearn.manifold import TSNE
    # # z是编码器生成的压缩向量.通过TSNE把z降成2维,画成散点图
    # normal_z = normal_z[:10000]
    # ano_z = ano_z[:10000]
    # z = tf.concat([normal_z, ano_z], axis=0)
    # digits_proj = TSNE().fit_transform(z)
    # plt.figure(figsize=(8, 8))
    # ax = plt.subplot(aspect='equal')
    # n_normal = normal_z.shape[0]
    # plt.suptitle("蓝色:正常样本,红色:异常样本")
    # ax.scatter(digits_proj[:n_normal, 0], digits_proj[:n_normal, 1], color='b')  # 正常样本用蓝色
    # ax.scatter(digits_proj[n_normal:, 0], digits_proj[n_normal:, 1], color='r')  # 异常样本用红色
    plt.show()


if __name__ == '__main__':
    app.run(main)

# python .\anomaly_detection\mae_dagmm.py  --corpus cwru --data_dim 1d --data_trans original
# python .\anomaly_detection\mae_dagmm.py  --corpus jiangnan --data_dim 1d --data_trans original